import random
import llm
import database_util
import time
from mylogger import logger

# model="Qwen/Qwen2.5-7B-Instruct"
model="01-ai/Yi-1.5-9B-Chat-16K"
def generate_my_story_QWenAI(storySize=1000, chapter_size=500):
    # ChatAI = llm.QWenChat()
    ChatAI = llm.SFChat(model=model)

    story = ""
    generateCount = 0

    GIRLS = ['翁建芳', '陈雅君', '张叶霞', '崔纯洁', '钱雅璐', '胡雅君', '戴丽芳', '黄雅敏', '张伶俐', '唐丽婷', '陈琳蕾', '陆茜', '谢丽娟', '万丝慧', '傅春霞', '张秀', '李霞', '沈雨婷', '黄佳燕', '仇芳萍', '卫丽娜', '张秋怡', '郭孟佳',]
    BOYS = ['沈卫飞', '邱万洪', '邓嘉飞', '李益君', '章佳俊', '蒋顺成', '怀岩鑫', '张磊', '郭豪', '郑晓斌', '王彦琪', '王陆鑫', '戚黎强', '翁剑浩', '朱天涛', '顾金松', '孙健成', '沈辉', '张家跃', '俞文彬', '李悦恒', '方叶勇', '章傲阳', '沈奂庆',]
    TEACHERS = ['李老师','陈老师','韩老师','马老师']
    STATES = [['2005', '一年'], ['2006', '二年'], ['2007', '三年']]
    SEASONS = ["春天", "夏天", "秋天", "冬天"]

    THEMES = ['青春成长', '校园爱情', '校园友情', '励志奋斗', '校园生活', '校园悬疑', '校园奇幻', '冒险', '攀登富士山', '攀登玉龙雪山', '泰国之旅', '长滩岛之旅', '攀登五姑娘山', '攀登长城',
        '惊悚', '校园科幻', '学霸故事', '学霸挑战学神', '穿越', '校园爽文', '科技革命', '时光穿越', '时间循环', '帮助老师', '获得全国性大奖', '获得世界级大奖', '暑假周游世界', '南极探险', '攀登珠峰',
        '海底寻宝', '拯救地球', '火星冒险', '实现常温超导', '探索月球', '探索地心', '驾驶高达', '不死药', '成长烦恼', '反乌托邦', '成长伤痛', '减肥', '参加校运会', '发现宝藏', '保护动物', '保护植物', '学霸成长记',
        '学渣成为学霸', '校园文艺汇演', "变成超人", "夺取全国冠军", '最强人工智能', '外星人', '发现亚特兰蒂斯' , '埃及探险', '法国游', '爱尔兰游' ,'冰岛游' , '探索马里亚纳海沟' ,'挖出国宝上交国家',
        '拯救濒危植物', '拯救濒危动物' , '发现陨石', '成为非洲酋长' , '彩票中头奖', '拯救全人类', '少年企业家' , '国际学生项目' , '老师无私奉献', '校园文化传承', '环保与社会责任', '脑机接口',
        '拍摄电影','参加世界杯','寻找泰坦尼克号', '编程大赛', '校园运动会', '南极大冒险', '西湖之旅', '武功山徒步', '穿越巴西', '穿越到古埃及', '穿越到古巴比伦', '穿越到中世纪'
    ]

    season = random.choice(SEASONS)
    states = random.choice(STATES)
    boys = random.choices(BOYS, k=random.randint(1,4))
    girl = random.choice(GIRLS)
    theme = random.choice(THEMES)

    prompt_s = f"""你是一位传奇故事作家，请创作一个多章节的初中校园传奇故事，故事发生在{2025}年{season}的{states[1]}4班中，主角是女生{girl}; 男生{','.join(boys)}；还有{random.choice(TEACHERS)}。故事的思考点如下
    故事的主题是关于{theme}的多章节故事。请思考以下几点:
    背景：故事发生在哪个学校？这个学校的特点和氛围是什么？故事的时间和地点是什么？
    主角：主角是谁？他们的性格、特点、目标和动机是什么？他们在故事中的角色是什么？
    冲突：故事中的冲突是什么？主角遇到了哪些困难和挑战？他们是如何克服的？
    情节：故事的情节是什么？包括开头、发展、高潮和结局等。
    角色：除了主角之外，还有哪些角色？他们的性格、特点、目标和动机是什么？他们在故事中的角色是什么？
    主题：故事的主题是什么？它与故事的情节、角色和背景有什么联系？故事传达了什么信息或价值观？
    风格：故事的风格是什么？是轻松幽默、悲伤感人、悬疑惊悚、奇幻神秘还是其他？
    结局：故事的结局是什么？它是如何解决冲突的？它是否与主题和情节相符？
    修改：故事还需要哪些修改和完善？语言、情节、角色、主题等方面是否已经做到最好？
    现在请你思考并写下故事大纲.
    """
    # 在写作时请注意：使用通俗易懂的语言，贴近少年视角，加入适量对话使语言生动，现在请你立即创作第一章。

    # print(prompt_s)
    ChatAI.startChat(prompt_s)
    prompt = f'在写作时请注意：使用通俗易懂的语言，贴近少年视角，加入适量对话使语言生动，现在请你直接创作第一章正文,字数控制在{min(storySize,chapter_size)}。'
    last = False
    for index in ['二','三','四','五','六','七','八','九','十','十一','十二','十三','十四','十五','十六','十七','十八','十九','二十']:
        try:
            assistantReplay = ChatAI.startChat(prompt)
        except Exception as e:
            # 如果已经生成了 50% 以上，就返回数据了
            if len(story) /storySize > 0.8:
                logger.info(f"生成异常，返回已生成的数据：{len(story)}")
                break
            else:
                raise e
        story += assistantReplay
        if last:
            break
        story += "\n\n"
        generateCount += 1
        if (len(story) / generateCount) * 0.9 > (storySize - len(story)) :
            # // 如果剩余生成字数小于平均每次生成字数 0.9，就不再生成
            prompt = f'请继续创作最终章正文,字数控制在{min(storySize - len(story),chapter_size)}'
            last=True
        else:
            prompt = f'请继续创作第{index}章正文,字数控制在{min(storySize - len(story),chapter_size)}'
            

    return story
                            

def main():
    while True:
        try:
            text = generate_my_story_QWenAI(5000,1000)
            count = 0
            llm_model = model
            version = "story_v2"
            # print(text)
            # return
            database_util.insert_data(text, count, llm_model, version)
        except Exception as e:
            logger.exception(e)
            time.sleep(10)
            

main()